package com.leo.boot.util;

import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.concurrent.atomic.AtomicInteger;

import static cn.hutool.core.text.StrPool.DASHED;
import static cn.hutool.core.util.ObjectUtil.defaultIfNull;
import static com.leo.boot.util.CommonUtils.uuid;
import static com.leo.boot.util.RequestContext.USER;
import static java.lang.System.currentTimeMillis;
import static java.util.Arrays.stream;

@Slf4j
@UtilityClass
public class TaskDecorators {

    private static final String USER_ID = "userId";

    private static final String TRACE_ID = "traceId";

    private static final String TASK = "task";

    private static final AtomicInteger counter = new AtomicInteger();

    public static TaskDecorator reduce(TaskDecorator... taskDecorators) {
        return stream(taskDecorators)
                .reduce((a, b) -> runnable -> a.decorate(b.decorate(runnable)))
                .orElse(a -> a);
    }

    public static TaskDecorator of() {
        return reduce(user(), trace(), stat());
    }

    public static TaskDecorator trace() {
        return runnable -> {
            String uid = defaultIfNull(MDC.get(USER_ID), USER.get().getUserId());
            String tid = defaultIfNull(MDC.get(TRACE_ID), uuid());
            return () -> {
                MDC.put(USER_ID, uid);
                MDC.put(TRACE_ID, tid);
                try {
                    runnable.run();
                } finally {
                    MDC.remove(USER_ID);
                    MDC.remove(TRACE_ID);
                }
            };
        };
    }

    public static TaskDecorator user() {
        return runnable -> {
            RequestContext.User user = USER.get();
            return () -> {
                USER.set(user);
                try {
                    runnable.run();
                } finally {
                    USER.remove();
                }
            };
        };
    }

    public static TaskDecorator stat() {
        return runnable -> {
            String taskId = defaultIfNull(MDC.get(TASK), TASK) + DASHED + counter.incrementAndGet();
            log.debug("Task {} submit", taskId);
            return () -> {
                long start = currentTimeMillis();
                log.debug("Task {} start", taskId);
                try {
                    runnable.run();
                } finally {
                    long end = currentTimeMillis();
                    log.debug("Task {} end, total time {} ms", taskId, end - start);
                }
            };
        };
    }
}